# -*- coding: utf-8 -*-
"""
Created on 2025/12/4 9:51

@author: zhulan
"""
import csv
from pathlib import Path

from loguru import logger

"""
水利实体匹配器 - 基于 Trie 树增强方案
任务负责人: @zhulanlanlan (珠兰)
"""

from pypinyin import lazy_pinyin
from typing import Optional, List, Dict
import Levenshtein
import time


class TrieNode:
    """Trie 树节点"""

    def __init__(self):
        self.children: Dict[str, 'TrieNode'] = {}
        self.is_end: bool = False
        self.entity_data: Optional[dict] = None  # 存储实体完整信息


class Trie:
    """Trie 树实现"""

    def __init__(self):
        self.root = TrieNode()

    def insert(self, word: str, entity_data: dict):
        """插入词条"""
        node = self.root
        for char in word:
            if char not in node.children:
                node.children[char] = TrieNode()
            node = node.children[char]
        node.is_end = True
        node.entity_data = entity_data

    def search(self, word: str) -> Optional[dict]:
        """精确搜索"""
        node = self.root
        for char in word:
            if char not in node.children:
                return None
            node = node.children[char]
        return node.entity_data if node.is_end else None

    def prefix_search(self, prefix: str, limit: int = 0) -> List[dict]:
        """
        前缀搜索 - 用于输入联想
        
        Args:
            prefix: 前缀字符串
            limit: 最大返回数量，0 表示不限制
        """
        node = self.root
        for char in prefix:
            if char not in node.children:
                return []
            node = node.children[char]

        # DFS 收集匹配结果（支持提前终止）
        results = []
        self._collect_all(node, results, limit)
        return results

    def prefix_search_first(self, prefix: str) -> Optional[dict]:
        """
        前缀搜索 - 只返回第一个匹配（高性能版本）
        
        找到第一个匹配后立即返回，不遍历整个子树
        """
        node = self.root
        for char in prefix:
            if char not in node.children:
                return None
            node = node.children[char]

        # 快速找到第一个匹配
        return self._find_first(node)

    def _find_first(self, node: TrieNode) -> Optional[dict]:
        """DFS 找到第一个匹配立即返回"""
        if node.is_end:
            return node.entity_data
        for child in node.children.values():
            result = self._find_first(child)
            if result:
                return result
        return None

    def _collect_all(self, node: TrieNode, results: List[dict], limit: int = 0):
        if node.is_end:
            results.append(node.entity_data)
            if limit and len(results) >= limit:
                return
        for child in node.children.values():
            if limit and len(results) >= limit:
                return
            self._collect_all(child, results, limit)


class HydroEntityMatcher:
    """
    水利实体匹配器

    支持功能:
    1. 汉字精确匹配
    2. 拼音匹配（解决同音错字）
    3. 拼音首字母匹配
    4. 别名映射
    5. 编辑距离容错
    """

    def __init__(self):
        self.trie_hanzi = Trie()  # 汉字 Trie
        self.trie_pinyin = Trie()  # 拼音全拼 Trie
        self.trie_shouzimu = Trie()  # 首字母 Trie
        self.alias_map: Dict[str, str] = {}  # 别名 → 标准名
        self.alias_pinyin_map: Dict[str, str] = {}  # 别名拼音 → 标准名（预计算，提升性能）
        self.all_entities: List[dict] = []  # 所有实体（用于模糊匹配）

    def add_entity(self, name: str, entity_type: str, code: str,
                   aliases: List[str] = None):
        """
        添加实体

        Args:
            name: 标准名称，如 "尼尔基水库"
            entity_type: 实体类型，如 "RESERVOIR"
            code: 实体编码，如 "RSV_NELJ"
            aliases: 别名列表，如 ["尼尔基", "尼尔基水电站"]
        """
        entity_data = {
            "name": name,
            "type": entity_type,
            "code": code
        }

        # 1. 汉字 Trie
        self.trie_hanzi.insert(name, entity_data)

        # 2. 拼音 Trie (去掉空格，全小写)
        pinyin = ''.join(lazy_pinyin(name))
        self.trie_pinyin.insert(pinyin, entity_data)

        # 3. 首字母 Trie
        shouzimu = ''.join([p[0] for p in lazy_pinyin(name)])
        self.trie_shouzimu.insert(shouzimu, entity_data)

        # 4. 别名映射
        if aliases:
            for alias in aliases:
                self.alias_map[alias] = name
                # 别名也加入汉字 Trie
                self.trie_hanzi.insert(alias, entity_data)
                # 别名拼音也加入（同时预存储到 alias_pinyin_map）
                alias_pinyin = ''.join(lazy_pinyin(alias))
                self.trie_pinyin.insert(alias_pinyin, entity_data)
                self.alias_pinyin_map[alias_pinyin] = name  # 预计算拼音映射

        # 5. 存储用于模糊匹配
        self.all_entities.append({
            **entity_data,
            "pinyin": pinyin,
            "shouzimu": shouzimu
        })

    def match(self, query: str, max_results: int = 5) -> List[dict]:
        """
        匹配实体

        匹配优先级:
        1. 别名表精确匹配
        2. 汉字 Trie 精确匹配
        3. 汉字 Trie 前缀匹配
        4. 拼音 Trie 匹配（解决同音错字）
        5. 首字母 Trie 匹配
        6. 编辑距离模糊匹配（兜底）

        Args:
            query: 用户输入
            max_results: 最大返回结果数

        Returns:
            匹配到的实体列表，按相关性排序
        """
        results = []
        query_lower = query.lower().strip()

        # 1. 别名表精确匹配
        if query in self.alias_map:
            standard_name = self.alias_map[query]
            result = self.trie_hanzi.search(standard_name)
            if result:
                results.append({**result, "match_type": "alias", "score": 1.0})

        # 2. 汉字 Trie 精确匹配
        if not results:
            result = self.trie_hanzi.search(query)
            if result:
                results.append({**result, "match_type": "exact", "score": 1.0})

        # 2.5 子串包含匹配（输入中包含实体名或别名）
        if not results:
            # 先做汉字子串匹配（快）
            for alias, standard_name in self.alias_map.items():
                if len(alias) >= 2 and alias in query:
                    result = self.trie_hanzi.search(standard_name)
                    if result:
                        score = len(alias) / len(query) * 0.95
                        results.append({**result, "match_type": "contains_alias", "score": score})
            
            # 汉字子串没匹配到，再做拼音子串匹配（使用预计算的拼音映射）
            if not results:
                query_pinyin = ''.join(lazy_pinyin(query))
                
                # 使用预计算的 alias_pinyin_map，避免重复计算拼音
                for alias_pinyin, standard_name in self.alias_pinyin_map.items():
                    if len(alias_pinyin) >= 3 and alias_pinyin in query_pinyin:
                        result = self.trie_hanzi.search(standard_name)
                        if result:
                            score = len(alias_pinyin) / len(query_pinyin) * 0.90
                            results.append({**result, "match_type": "contains_alias_pinyin", "score": score})
                
                # 检查实体标准名的拼音
                if not results:
                    for entity in self.all_entities:
                        name_pinyin = entity["pinyin"]
                        if len(name_pinyin) >= 4 and name_pinyin in query_pinyin:
                            score = len(name_pinyin) / len(query_pinyin) * 0.85
                            results.append({
                                "name": entity["name"],
                                "type": entity["type"],
                                "code": entity["code"],
                                "match_type": "contains_name_pinyin",
                                "score": score
                            })
            
            # 去重并按分数排序
            if results:
                seen = set()
                unique = []
                for r in sorted(results, key=lambda x: x["score"], reverse=True):
                    if r["name"] not in seen:
                        seen.add(r["name"])
                        unique.append(r)
                results = unique[:max_results]

        # 3. 汉字 Trie 前缀匹配
        if not results:
            prefix_results = self.trie_hanzi.prefix_search(query)
            for r in prefix_results[:max_results]:
                results.append({**r, "match_type": "prefix", "score": 0.9})

        # 4. 拼音 Trie 匹配
        if not results:
            query_pinyin = ''.join(lazy_pinyin(query))
            result = self.trie_pinyin.search(query_pinyin)
            if result:
                # 拼音精确匹配，分数应该很高（仅次于汉字精确匹配）
                results.append({**result, "match_type": "pinyin", "score": 0.95})
            else:
                # 拼音前缀匹配
                pinyin_results = self.trie_pinyin.prefix_search(query_pinyin)
                for r in pinyin_results[:max_results]:
                    results.append({**r, "match_type": "pinyin_prefix", "score": 0.85})

        # 5. 首字母匹配
        if not results and query_lower.isalpha():
            result = self.trie_shouzimu.search(query_lower)
            if result:
                results.append({**result, "match_type": "shouzimu", "score": 0.80})

        # 6. 编辑距离模糊匹配（兜底）
        if not results:
            query_pinyin = ''.join(lazy_pinyin(query))
            fuzzy_results = []
            for entity in self.all_entities:
                # 计算拼音编辑距离
                distance = Levenshtein.distance(query_pinyin, entity["pinyin"])
                if distance <= 3:  # 允许最多3个字符差异
                    score = 1 - (distance / max(len(query_pinyin), len(entity["pinyin"])))
                    fuzzy_results.append({
                        "name": entity["name"],
                        "type": entity["type"],
                        "code": entity["code"],
                        "match_type": "fuzzy",
                        "score": score * 0.7,
                        "distance": distance
                    })

            # 按分数排序
            fuzzy_results.sort(key=lambda x: x["score"], reverse=True)
            results.extend(fuzzy_results[:max_results])

        return results[:max_results]

    def load_from_csv(self, csv_path: str, encoding: str = 'utf-8') -> int:
        """
        从 CSV 文件加载实体数据

        CSV 格式要求:
        - 列: code, name, type, aliases
        - aliases 使用 | 分隔多个别名

        Args:
            csv_path: CSV 文件路径
            encoding: 文件编码，默认 utf-8

        Returns:
            加载的实体数量
        """
        count = 0
        path = Path(csv_path)

        if not path.exists():
            logger.error(f"CSV 文件不存在: {csv_path}")
            return 0

        with open(path, 'r', encoding=encoding) as f:
            reader = csv.DictReader(f)
            for row in reader:
                code = row.get('code', '').strip()
                name = row.get('name', '').strip()
                entity_type = row.get('type', '').strip()
                aliases_str = row.get('aliases', '').strip()

                if not name or not entity_type or not code:
                    logger.warning(f"跳过无效行: {row}")
                    continue

                # 解析别名（使用 | 分隔）
                aliases = [a.strip() for a in aliases_str.split('|') if a.strip()] if aliases_str else None

                self.add_entity(name, entity_type, code, aliases)
                count += 1

        logger.info(f"从 CSV 加载 {count} 个实体: {csv_path}")
        return count


# ========== 使用示例 ==========

if __name__ == "__main__":
    # 初始化匹配器
    matcher = HydroEntityMatcher()

    # # 添加实体数据
    # matcher.add_entity(
    #     name="尼尔基水库",
    #     entity_type="RESERVOIR",
    #     code="RSV_NELJ",
    #     aliases=["尼尔基", "尼尔基水电站"]
    # )
    # matcher.add_entity(
    #     name="丰满水库",
    #     entity_type="RESERVOIR",
    #     code="RSV_FM",
    #     aliases=["丰满", "丰满大坝", "丰满电站"]
    # )
    # matcher.add_entity(
    #     name="第二松花江",
    #     entity_type="RIVER",
    #     code="RIV_ESHJ",
    #     aliases=["二松", "二松花江"]
    # )
    # matcher.add_entity(
    #     name="江桥水文站",
    #     entity_type="STATION",
    #     code="STN_JQ",
    #     aliases=["江桥站", "江桥"]
    # )
    # 从csv加载数据到matcher
    csv_path = Path(__file__).parent.parent / "data" / "hydro_entities.csv"
    matcher.load_from_csv(str(csv_path))

    # 测试各种匹配场景
    test_cases = [
        "碧流河水库",    # 精确匹配
        "碧流河",      # 别名匹配
        "碧流",        # 前缀匹配
        "碧柳河水库",    # 同音错字 → 拼音匹配
        "必留河",      # 同音错字 → 拼音匹配
        "大伙房",    # 别名匹配
        "逼瘤鹤",  # 首字母匹配 (碧流河水库)
        "必liu河",  # 首字母匹配 (碧流河水库)
        "blh",      # 首字母匹配 (碧流河水库)
        "blhsk",      # 首字母匹配 (碧流河水库)
        "dhfsk",      # 首字母匹配 (大伙房水库)
        "biliuhe",    # 拼音匹配
        "碧流和水库",    # 同音错字
        "碧流水库",     # 漏字 → 模糊匹配
        "清河水库",     # 精确匹配
        "清河",        # 别名/前缀匹配
        "大伙房傻逼",
        "傻逼大伙房",
        "小战马在blh狂草zjx",
        "大王在碧流河吃小战马鸡巴",
        "大王在逼流河吃小战马鸡巴"
    ]

    logger.info("=" * 60)
    logger.info("水利实体匹配测试")
    logger.info("=" * 60)
    t = time.perf_counter()

    for query in test_cases:
        t0 = time.perf_counter()
        results = matcher.match(query)
        elapsed_ms = (time.perf_counter() - t0) * 1000
        # results按name去重，保留第一个（分数最高的）
        seen = set()
        unique_results = []
        for r in results:
            if r['name'] not in seen:
                seen.add(r['name'])
                unique_results.append(r)
        results = unique_results
        if results:
            for result in results:
                logger.debug(f"输入: {query:12} → 匹配: {result['name']:10} "
                             f"(方式: {result['match_type']}, 分数: {result['score']:.2f}, 耗时: {elapsed_ms:.3f} ms)")
            logger.success(f"查询 {query} 的匹配 {results}")
        else:
            logger.error(f"输入: {query:12} → 未匹配 (耗时: {elapsed_ms:.3f} ms)")

    logger.info("=" * 60)
    elapsed = (time.perf_counter() - t) * 1000
    logger.debug(f"输入: {len(test_cases)} 问题 "
                 f"(耗时: {elapsed:.3f} ms)")
